import argparse
from sys import platform
from models import *
from utils.datasets import *
from utils.utils import *
import pickle

import numpy as np
np.set_printoptions(threshold=10000)

def detect(save_txt=False, save_img=False):
    img_size = opt.img_size  # (320, 192) or (416, 256) or (608, 352) for (height, width)
    out, source, weights= opt.output, opt.source, opt.weights

    # Initialize
    device = torch_utils.select_device(opt.device)
    if os.path.exists(out):
        shutil.rmtree(out)  # delete output folder
    os.makedirs(out)  # make new output folder

    # Initialize model
    model = Darknet(opt.cfg, img_size, quan=opt.quan, save_act=opt.save_act, conv_bias=not opt.no_bias)

    # Load weights
    if weights.endswith('.pt'):  # pytorch format
        model.load_state_dict(torch.load(weights, map_location=device)['model'])
    
    # Eval mode
    model.to(device).eval()
    dataset = LoadImages(source, img_size=img_size, half=False)

    # Get names and colors
    names = load_classes(opt.names)
    colors = [[random.randint(0, 255) for _ in range(3)] for _ in range(len(names))]
    
    if opt.prune:
        prune_weight(model, opt.threshold)
    
    if opt.quan:
        quan_weight(model)
    
    # Run inference
    t0 = time.time()
    # outfile=dict()
    for path, img, im0s, vid_cap in dataset:
        #t = time.time()
        ########
        # with open('./output/input.txt', 'wb') as f:
        #     pickle.dump(img, f)
        ########    
        # Get detections
        img = torch.from_numpy(img).to(device)
        
            
        if img.ndimension() == 3:
            img = img.unsqueeze(0)
        t = time.time()
        pred = model(img)[0]

        t2 = time.time()
        print('time = ',t2 - t)
        print('pred.shape=', pred.shape)
        # Apply NMS
        pred = non_max_suppression(pred, opt.conf_thres, opt.nms_thres)
        t3 = time.time()
        print('time = ',t3 - t2)

        print('pred=', pred)
        # Process detections
        for i, det in enumerate(pred):  # detections per image
            p, s, im0 = path, '', im0s

            save_path = str(Path(out) / Path(p).name)
            s += '%gx%g ' % img.shape[2:]  # print string
            if det is not None and len(det):
                # Rescale boxes from img_size to im0 size
                det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()

                # Print results
                for c in det[:, -1].unique():
                    n = (det[:, -1] == c).sum()  # detections per class
                    s += '%g %ss, ' % (n, names[int(c)])  # add to string

                # Print time (inference + NMS)
                print('%sDone. (%.3fs)' % (s, time.time() - t))

                # Write results
                for *xyxy, conf, cls in det:
                    label = '%s %.2f' % (names[int(cls)], conf)
                    plot_one_box(xyxy, im0, label=label, color=colors[int(cls)])

            # Save results (image with detections)
            cv2.imwrite(save_path, im0)
        
        if opt.save_act:
            model.save_act_pkl(opt.out_act)

        '''
        with open(opt.out_act, 'rb') as f:
            act=pickle.load(f)
        
        act_new=dict()
        i=0
        for key in act.keys():
            if 'convolutional' in key:
                name='convolutional'+str(i)
                i+=1
                act_new[name]=act[key]

        outfile[path]={}
        for key in act_new.keys():
            z=int(torch.sum(act_new[key]==0))
            s=act_new[key].shape
            m=1
            for item in s:
                m*=item 
            outfile[path][key]=z/m
        '''
    if opt.save_wgt:
        wgt_parm={}
        for name,parameters in model.named_parameters():
            if 'Conv2d' in name:
                wgt_parm[name]=parameters.cpu().detach().numpy()        
        import pickle
        f_wgt=open(opt.out_wgt,'wb')
        pickle.dump(wgt_parm, f_wgt)
        f_wgt.close()

    '''
    with open('./20200730.csv','a') as f:
        for pic_path in outfile.keys():
            f.write(str(pic_path))
            pic_result = outfile[pic_path]
            for l in pic_result.keys():
                f.write(', ')
                f.write(str(pic_result[l]))
                
            f.write('\n')
   
    parm={}
    for name,parameters in model.named_parameters():
        parm[name]=parameters.cpu().detach().numpy()
    zero_num=0
    weight_num=0
    for key in parm.keys():
        if(('scale' not in key) and ('bias' not in key)):
            s=parm[key].shape
            m=1
            for item in s:
                m*=item
            z=np.sum(parm[key]==0)

            weight_num+=m
            zero_num+=z
    print('---------------------')
    print('zero_num=',zero_num)
    print('weight_num=',weight_num)
    print('zero_num/weight_num=',zero_num/weight_num)    
    '''
    print('Results saved to %s' % os.getcwd() + os.sep + out)
    print('Done. (%.3fs)' % (time.time() - t0))
    
if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--cfg', type=str, default='cfg/yolov3_ReLU_fuse.cfg', help='*.cfg path')
    parser.add_argument('--names', type=str, default='data/coco.names', help='*.names path')
    parser.add_argument('--weights', type=str, default='/home/zhangjm/backup/20201120/yolov3_ReLU_fuse_no_bias.pt', help='path to weights file')
    parser.add_argument('--source', type=str, default='/home/zhangjm/yolov3_quan/data/samples/destop_new.jpg', help='source')  # input file/folder, 0 for webcam
    parser.add_argument('--output', type=str, default='output', help='output folder')  # output folder
    parser.add_argument('--img-size', type=int, default=416, help='inference size (pixels)')
    parser.add_argument('--conf-thres', type=float, default=0.3, help='object confidence threshold')
    parser.add_argument('--nms-thres', type=float, default=0.5, help='iou threshold for non-maximum suppression')

    parser.add_argument('--device', default='', help='device id (i.e. 0 or 0,1) or cpu')
    parser.add_argument('--quan', action='store_true', help='quantize the model')
    parser.add_argument('--no_bias', action='store_true', help='whether model has bias')
    parser.add_argument('--save-act', action='store_true', help='save activation')
    parser.add_argument('--save-wgt', action='store_true', help='save wgt')
    parser.add_argument('--out-act', type=str, default='/home/zhangjm/backup/20210310/destop_new.pkl', help='save activation')
    parser.add_argument('--out-wgt', type=str, default='/home/zhangjm/backup/20210310/1923wgt.pkl', help='save wgt')
    parser.add_argument('--prune', action='store_true', help='prune the model')
    parser.add_argument('--threshold', type=float, default=0.005)
    
    opt = parser.parse_args()
    print(opt)

    with torch.no_grad():
        detect()
